#include "llvm/ADT/DenseMap.h"
#include "llvm/ADT/DepthFirstIterator.h"
#include "llvm/ADT/GraphTraits.h"
#include "llvm/ADT/MapVector.h"
#include "llvm/ADT/PostOrderIterator.h"
#include "llvm/ADT/SetVector.h"
#include "llvm/ADT/SmallPtrSet.h"
#include "llvm/ADT/SmallVector.h"
#include "llvm/ADT/Statistic.h"
#include "llvm/Analysis/DomTreeUpdater.h"
#include "llvm/Analysis/GlobalsModRef.h"
#include "llvm/Analysis/IteratedDominanceFrontier.h"
#include "llvm/Analysis/PostDominators.h"
#include "llvm/IR/BasicBlock.h"
#include "llvm/IR/CFG.h"
#include "llvm/IR/DebugInfoMetadata.h"
#include "llvm/IR/DebugLoc.h"
#include "llvm/IR/Dominators.h"
#include "llvm/IR/Function.h"
#include "llvm/IR/IRBuilder.h"
#include "llvm/IR/InstIterator.h"
#include "llvm/IR/InstrTypes.h"
#include "llvm/IR/Instruction.h"
#include "llvm/IR/Instructions.h"
#include "llvm/IR/IntrinsicInst.h"
#include "llvm/IR/PassManager.h"
#include "llvm/IR/Use.h"
#include "llvm/IR/Value.h"
#include "llvm/InitializePasses.h"
#include "llvm/Pass.h"
#include "llvm/ProfileData/InstrProf.h"
#include "llvm/Support/Casting.h"
#include "llvm/Support/CommandLine.h"
#include "llvm/Support/Debug.h"
#include "llvm/Support/raw_ostream.h"
#include "llvm/Transforms/Scalar.h"

#include <iostream>
#include <vector>
#include <map>
#include <string>
#include <unistd.h>
#include <string.h>
#include <utility>      // std::pair, std::make_pair
#include <algorithm>
#include <memory>       // std::shared_ptr

using namespace llvm;

namespace llvm {
    FunctionPass * createLoopCountPass();
    void initializeLoopCountPassPass(PassRegistry&);
}

namespace { //在这个namespace下,写要用到的数据结构，函数，类

    template <class Ptr>
    class BitSet{
    public:
        char* bits;	//根据alphabet的大小动态分配
        BitSet(std::map<Ptr,int> *_alphabet):alphabet(_alphabet){
            frame_size = sizeof(char)*8;
            size = alphabet->size();
            length = (size+frame_size-1)/frame_size;
            bits = (char *)malloc(length*frame_size);
            memset(bits,0,length*frame_size);
        }
        ~BitSet(){
            delete[] bits;
        }
        void insert(Ptr elem){
            int idx = (*alphabet)[elem];
            bits[idx/frame_size] |= 1<<(idx%frame_size);
        }
        bool has(Ptr elem){
            int idx = (*alphabet)[elem];
            return bits[idx/frame_size] & (1<<(idx%frame_size));
        }
        void join(std::shared_ptr< BitSet<Ptr> > s){
            //假定他们的全集是一样的
            for(int i = 0; i < length; ++i){
                bits[i] |= s->bits[i];
            }
        }
        bool isSubsetOf(std::shared_ptr< BitSet<Ptr> > s){
            //假定他们的全集是一样的
            for(int i = 0; i < length; ++i){
                if((bits[i] & s->bits[i]) != bits[i]){
                    return false;
                }
            }
            return true;
        }
        void print(){
            for(int i = 0; i < length; ++i){
                for(int j = 0; j < frame_size; ++j){
                    if(bits[i]&(1<<j)){
                        std::cout << "1";
                    }
                    else{
                        std::cout << "0";
                    }
                }
            }
            std::cout << std::endl;
        }
        void full(){
            //将集合变成全集
            memset(bits,0xffffffff,length*frame_size);
        }
    private:
        std::map<Ptr,int> *alphabet;
        int frame_size;
        int size;	//等于alphabet->size()
        int length;	//等于bits所指的数组的大小
    };

    typedef BitSet<BasicBlock*> LoopSet;
    typedef std::shared_ptr<LoopSet> PtrLoopSet;
    typedef std::pair<BasicBlock*,PtrLoopSet> LoopRecord;

    class Node;
    typedef std::shared_ptr<Node> PtrNode;

    class Node{ //表示循环嵌套关系的树的节点
    public:
        std::vector<PtrNode> children;
        PtrNode father;
        PtrLoopSet loop;
        BasicBlock *header;
        Node(PtrNode _father, PtrLoopSet _loop, BasicBlock *_header):father(_father),loop(_loop),header(_header){}
        inline bool isRoot(){
            return father == nullptr;
        }
        inline void addChild(PtrNode child){
            children.push_back(child);
        }
    };

    std::map<BasicBlock*,int> alphabet;
    std::vector<LoopRecord> loops;
    

//...写需要用到的函数，在runOnFunction方法中调用
    void getAlphabet(BasicBlock *p, int &dfs_idx){
        //在CFG上递归地dfs，建立BasicBlock*到dfs_idx的映射关系，保存到alphabet中
        if(alphabet.count(p))
            return;
        alphabet[p] = dfs_idx++;
        for(auto I = succ_begin(p); I != succ_end(p); ++I){
            getAlphabet(*I,dfs_idx);   //*I为指向后继基本块的指针
        }
    }

    void findLoop(BasicBlock *from, BasicBlock *to, PtrLoopSet loop){
        //找到回边from->to对应的循环，结果放到loop中
        loop->insert(from);
        for(auto I = pred_begin(from); I !=pred_end(from); ++I){
            if(!loop->has(*I)){
                findLoop(*I,to,loop);
            }
        }
    }

    void findBackEdge( DominatorTree &dom_tree){
        //在CFG上递归地dfs，找到所有的回边及其对应的循环，将循环保存到loops中
        std::vector<DomTreeNode*> WorkList;
        WorkList.push_back(dom_tree.getRootNode());
        while(!WorkList.empty()){
            auto CurNode = WorkList.back();
            WorkList.pop_back();
            auto BB = CurNode->getBlock();
            for(auto sSucc = succ_begin(BB), eSucc = succ_end(BB);sSucc != eSucc; ++sSucc){
                auto SuccNode = dom_tree.getNode(*sSucc);
                if(dom_tree.dominates(SuccNode, CurNode)){
                    BasicBlock *from = CurNode->getBlock();
                    BasicBlock *to = SuccNode->getBlock();
                    PtrLoopSet loop(new LoopSet(&alphabet));
                    loop->insert(to);
                    findLoop(from,to,loop);
                    loops.push_back(std::pair<BasicBlock*,PtrLoopSet>(to,loop));
                }
            }
            WorkList.insert(WorkList.end(), CurNode->begin(), CurNode->end());
        }
    }

    void findAllLoop(Function &F){
        //dfs整个CFG，找到所有的回边，并对每个回边找出其对应的自然循环
        //dfs序从0开始记，每个BB对应的dfs序存入alphabet中
        //loops中保存所有的header和对应的循环
        BasicBlock *entry = &F.getEntryBlock();
        if (succ_empty(entry))
            return;
        int dfs_idx = 0;
        getAlphabet(entry, dfs_idx);
        DominatorTree dom_tree(F);  //按照LoopSearchPass，构造函数应传入F而不是&F

        findBackEdge(dom_tree);    //会同时计算每个回边对应的循环
    }
    
    bool cmp_loop(LoopRecord &a, LoopRecord &b){
        return alphabet[a.first] < alphabet[b.first];
    }

    PtrNode buildLoopTree(){
        //建立表示循环嵌套关系的树，返回值是树的根节点
        std::sort(loops.begin(), loops.end(), cmp_loop);
        PtrLoopSet uni_set(new LoopSet(&alphabet));
        uni_set->full();     //变成全集
        PtrNode root(new Node(NULL,uni_set,NULL));
        PtrNode p = root;
        for(auto loop_record : loops){
            BasicBlock *header = loop_record.first;
            PtrLoopSet loop = loop_record.second;
            if(alphabet[header] == alphabet[p->header]){
                p->loop->join(loop);
                continue;
            }
            while(!(loop->isSubsetOf(p->loop))){
                p = p->father;
            }
            PtrNode new_node(new Node(p,loop,header));
            p->addChild(new_node);
            p = new_node;
        }
        return root;
    }

    void printTree(std::ostream &os, PtrNode p, const std::string &prefix, int idx, int depth){
        //递归打印整个树
        std::string name;
        if(p->isRoot())
            name = prefix;
        else
            name = prefix + std::to_string(idx);
        
        std::string indent(depth*4, ' ');

        if(!(p->isRoot())){
            os << indent << "\"" << name << "\" : ";
        }
        os << "{\n";
        if(!(p->isRoot())){
            os << indent << "    \"depth\" : " << depth;
            if(p->children.size())
                os << ",";
            os << "\n";
        }
        for(int i = 0; i < p->children.size(); i++){
            printTree(os, p->children[i], name, i+1, depth+1);
            if(i != p->children.size() - 1){
                os << ",";
            }
            os << "\n";
        }
        os << indent << "}";
        if(p->isRoot()){
            os << "\n";
        }
    }

    struct LoopCountPass : public FunctionPass {
        static char ID;
        LoopCountPass() : FunctionPass(ID) {
            initializeLoopCountPassPass(*PassRegistry::getPassRegistry());
        }
        bool runOnFunction(Function &F) override {  //运行后，driver会调用FunctionPass的runOnFunction
            if (skipFunction(F))
                return false;
            findAllLoop(F);
            PtrNode root = buildLoopTree();
            printTree(std::cout,root,"L",-1,0);
            return true;
        }
        void getAnalysisUsage(AnalysisUsage &AU) const override {
            AU.addRequired<PostDominatorTreeWrapperPass>();
            AU.setPreservesCFG();
            AU.addPreserved<DominatorTreeWrapperPass>();
            AU.addPreserved<PostDominatorTreeWrapperPass>();
            AU.addPreserved<GlobalsAAWrapperPass>();
        }
    };
}

char LoopCountPass::ID = 0;
INITIALIZE_PASS(LoopCountPass, "loopcount", "Loop Message Count", false, false)
FunctionPass *llvm::createLoopCountPass() {
    return new LoopCountPass();
}
